# -*- coding: utf-8 -*-
from __future__ import division
import numpy as np
import scipy.io as sio
import matplotlib.pyplot as plt
from scipy.signal import butter, cheby2, filtfilt


def bandpass_filter(data_x, Fs, filter_low, filter_high, filter_bank=False, filter_num=1):
    """
    Chebyshev type II filter

    输入参数
    ----------
    data_x: T×N×L ndarray(或单个trial T×N)
           T: 采样点数  N: 通道数  L: 训练数据trial总数
       Fs: 采样频率

    返回值
    ----------
    after_filter_x: T×N×L(或单个trial T×N)

    """
    channel_num = data_x.shape[1]
    filter_order = 6
    rs = 40
    if filter_bank:
        filter_width = (filter_high - filter_low) / filter_num
        filter_range = [filter_low + filter_width * i for i in range(filter_num + 1)]
        wn = [i / (Fs / 2) for i in filter_range]
        wn1 = wn[0:filter_num]
        wn2 = wn[1:]
        if len(data_x.shape) == 3:  # 输入为多个trial时
            trial_size = data_x.shape[2]
            after_filter_x = np.zeros([data_x.shape[0], channel_num, filter_num, trial_size])
            for k in range(filter_num):
                [b, a] = cheby2(filter_order, rs, Wn=[wn1[k], wn2[k]], btype='bandpass')
                for i in range(trial_size):
                    for j in range(channel_num):
                        after_filter_x[:, j, k, i] = filtfilt(b, a, data_x[:, j, i])
        else:  # 输入为单个trial
            after_filter_x = np.zeros([data_x.shape[0], channel_num, filter_num])
            for k in range(filter_num):
                [b, a] = cheby2(filter_order, rs, Wn=[wn1[k], wn2[k]], btype='bandpass')
                for j in range(channel_num):
                    after_filter_x[:, j, k] = filtfilt(b, a, data_x[:, j])
    else:
        wn1 = filter_low / (Fs / 2)
        wn2 = filter_high / (Fs / 2)
        if len(data_x.shape) == 3:  # 输入为多个trial时
            trial_size = data_x.shape[2]
            after_filter_x = np.zeros(data_x.shape)
            [b, a] = cheby2(filter_order, rs, Wn=[wn1, wn2], btype='bandpass')
            for i in range(trial_size):
                for j in range(channel_num):
                    after_filter_x[:, j, i] = filtfilt(b, a, data_x[:, j, i])
        else:  # 输入为单个trial
            after_filter_x = np.zeros(data_x.shape)
            [b, a] = cheby2(filter_order, rs, Wn=[wn1, wn2], btype='bandpass')
            for j in range(channel_num):
                after_filter_x[:, j] = filtfilt(b, a, data_x[:, j])
    return after_filter_x


if __name__ == '__main__':
    from scipy.fftpack import fft
    data_xForFilter = sio.loadmat(r'D:\Myfiles\MI-BCI源码\PythonProjects\Data\trainx.mat')
    data_x = data_xForFilter['train_x']  # shape(750,22,138)
    Fs = 500
    filter_low = 4
    filter_high = 40
    t = np.linspace(1, data_x.shape[0], data_x.shape[0])
    data = data_x[:, :, 2]
    AfterFilter_x = bandpass_filter(data, Fs, filter_low, filter_high)
    data = data[:, 1]
    AfterFilter_x = AfterFilter_x[:, 1]
    x_fft = abs(fft(data))
    AfterFilter_x_fft = abs(fft(AfterFilter_x))
    # x_fft = x_fft[range(int(len(data)/2))]
    # AfterFilter_x_fft = AfterFilter_x_fft[range(int(len(data)/2))]
    plt.subplot(221)
    plt.plot(t, data)
    plt.subplot(222)
    plt.plot(t, AfterFilter_x)
    plt.subplot(223)
    plt.plot(t, x_fft)
    plt.subplot(224)
    plt.plot(t, AfterFilter_x_fft)
    plt.show()
